# Credit: This file generates LLVM IR from ISPCTextureCompressor from https://github.com/GameTechDev/ISPCTextureCompressor/tree/master

import os
import requests
import subprocess
from shutil import which

REPO_URL = "https://github.com/GameTechDev/ISPCTextureCompressor"
BC6H_KERNEL = "CompressBlocksBC6H_ispc"
BC7_KERNEL = "CompressBlocksBC7_ispc"


def download_ispc_tex_compress(output_file: str):
    url = f"{REPO_URL}/raw/refs/heads/master/ispc_texcomp/kernel.ispc"
    src = requests.get(url).content.decode("utf-8")
    with open(output_file, "w") as f:
        f.write(src)


def compile_ispc_kernel(src_file: str, output_file: str, arch: str, ispc_exe: str):
    if arch == "arm64":
        subprocess.run([ispc_exe, src_file, "-o", output_file, "-O3", "-DNDEBUG", "--emit-llvm-text",
                        "--opt=disable-assertions", "--opt=fast-math", "--math-lib=fast", "-woff",
                        "--opt=enable-ldst-vectorizer", "--opt=enable-slp-vectorizer", f"--target=neon-i32x8"])
    else:
        subprocess.run([ispc_exe, src_file, "-o", output_file, "-O3", "-DNDEBUG", "--emit-llvm-text",
                        "--opt=disable-assertions", "--opt=fast-math", "--math-lib=fast", "-woff"])


def patch_llvm_ir(src_file: str, output_file: str, preserved_functions: list):
    with open(src_file, "r") as f:
        src = f.readlines()
    patched_lines = [line.replace("define ", "define private ")
                     if (line.startswith("define ") and
                         not any(line.startswith(f"define void @{f}(") for f in preserved_functions))
                     else line for line in src]
    with open(output_file, "w") as f:
        f.writelines(patched_lines)


def compile_llvm_ir(src_file: str, output_file: str, clang_exe: str):
    subprocess.run([clang_exe, src_file, "-o", output_file, "-O3", "-emit-llvm", "-c"])


def embed_bitcode_into_c_hex_array(bc_file: str, output_file: str, array_name: str):
    with open(bc_file, "rb") as f:
        bc = f.read()
    with open(output_file, "w") as f:
        f.write(f"// This file is generated by generate_ispc_tex_compress.py\n")
        f.write(f"// See original ISPC kernel: {REPO_URL}\n\n")
        f.write(f"static constexpr auto {array_name}_bc6h_kernel_name = \"{BC6H_KERNEL}\";\n")
        f.write(f"static constexpr auto {array_name}_bc7_kernel_name = \"{BC7_KERNEL}\";\n\n")
        f.write(f"static const uint8_t {array_name}[] = {{\n")
        for i in range(0, len(bc), 16):
            f.write("    " + ", ".join(f"0x{c:02x}u" for c in bc[i:i + 16]) + ",\n")
        f.write("};\n")


if __name__ == "__main__":
    base_dir = os.path.dirname(os.path.realpath(__file__))
    # download the tex_compress.ispc file
    src_path = os.path.join(base_dir, "tex_compress.ispc")
    download_ispc_tex_compress(src_path)
    ispc_exe = which("ispc")
    if ispc_exe is None:
        raise FileNotFoundError("ispc not found in PATH")
    clang_exe = which("clang")
    if clang_exe is None:
        raise FileNotFoundError("clang not found in PATH")
    # compile the ispc kernel to LLVM IR
    arch = "x86_64" if os.uname().machine == "x86_64" else "arm64"
    ll_path = os.path.join(base_dir, "tex_compress.ll")
    compile_ispc_kernel(src_path, ll_path, arch, ispc_exe)
    # patch the LLVM IR to make unwanted functions private
    patched_ll_path = os.path.join(base_dir, "tex_compress.patched.ll")
    patch_llvm_ir(ll_path, patched_ll_path, [BC6H_KERNEL, BC7_KERNEL])
    # compile the LLVM IR to bitcode
    bc_path = os.path.join(base_dir, "tex_compress.bc")
    compile_llvm_ir(patched_ll_path, bc_path, clang_exe)
    # embed the bitcode into C hex array in "fallback_tex_compress.inl"
    c_path = os.path.join(base_dir, f"fallback_tex_compress.{arch}.inl")
    embed_bitcode_into_c_hex_array(bc_path, c_path, "fallback_tex_compress_llvm_bc")
